import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from skfem import *
from skfem.models import laplace, mass
from scipy.sparse import bmat, identity, csr_matrix
from scipy.sparse.linalg import splu

# 参数设置
L = 1.0          # 弦的长度
z = 0.3          # 施加力的位置（距离左端）
Fy = 0.01        # 施加的垂直力大小
T = 10.0         # 弦的张力
rho = 0.1        # 弦的线密度
c = np.sqrt(T/rho)  # 波速
duration = 0.5   # 模拟时长
dt = 0.001       # 时间步长

# 创建网格和基函数
mesh = MeshLine(np.linspace(0, L, 100)).with_boundaries({
    'left': lambda x: x[0] == 0,
    'right': lambda x: x[0] == L
})
basis = Basis(mesh, ElementLineP1())

# 组装质量矩阵和刚度矩阵
M = mass.assemble(basis)
K = laplace.assemble(basis) * T

# 获取边界自由度 - 使用不同的方法
boundary_dofs = []
for facet in ['left', 'right']:
    dofs = basis.get_dofs(facet)
    # 使用 DofsView 的 all() 方法获取所有自由度
    boundary_dofs.extend(dofs.all())

# 手动应用边界条件
def apply_bc(A):
    """应用边界条件到矩阵A"""
    A = A.copy().tocsc()
    for dof in boundary_dofs:
        # 将边界点的行和列清零，对角线设为1
        A[dof, :] = 0
        A[:, dof] = 0
        A[dof, dof] = 1
    return A

# 应用边界条件
M_bc = apply_bc(M)
K_bc = apply_bc(K)

# 初始条件：高斯形状的初始位移
sigma = 0.05
u0 = basis.project(lambda x: Fy * np.exp(-(x[0] - z)**2 / (2 * sigma**2)))

# 应用边界条件到初始位移
for dof in boundary_dofs:
    u0[dof] = 0

# 转换为状态空间形式（一阶系统）
N = M_bc.shape[0]  # 自由度数量
I = identity(N)

# 构建系统矩阵
A0 = bmat([[I, None], [None, M_bc]], 'csr')
B0 = bmat([[None, I], [-c**2 * K_bc, None]], 'csr')

# Crank-Nicolson方法
theta = 0.5
A = A0 + theta * B0 * dt
B = A0 - (1 - theta) * B0 * dt
backsolve = splu(A).solve

# 初始状态
U = np.concatenate([u0, np.zeros(N)])

# 时间步进函数
def evolve(t, u):
    while t < duration:
        t += dt
        u = backsolve(B.dot(u))
        yield t, u

# 创建动画
fig, ax = plt.subplots(figsize=(10, 6))
ax.set(xlim=(0, L), ylim=(-1.5*Fy, 1.5*Fy), 
       xlabel='Position x', ylabel='Displacement u(x,t)',
       title='String Vibration')
ax.grid(True)

# 按x坐标排序以便绘图
ix = np.argsort(mesh.p[0])
line, = ax.plot(mesh.p[0, ix], u0[ix], lw=2)
point, = ax.plot([z], [0], 'ro', ms=8)
time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes)

def update(t_u):
    t, u = t_u
    u1, _ = np.split(u, [N])
    line.set_ydata(u1[ix])
    time_text.set_text(f'Time = {t:.2f} s')
    return line, point, time_text

# 创建动画
animation = FuncAnimation(fig, update, evolve(0, U), interval=50)
plt.show()

# 保存动画
animation.save('string_vibration.gif', writer='pillow', fps=20)